package com.tool4j.node;

import cn.hutool.core.collection.CollUtil;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.tool4j.entity.TableMetadata;
import com.tool4j.util.SqlParserUtil;
import net.sf.jsqlparser.JSQLParserException;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.core.ParameterizedTypeReference;

import java.util.List;
import java.util.Map;

public class SqlParserNode implements NodeAction {

    private ChatClient chatClient;


    public SqlParserNode(ChatClient chatClient) {
        this.chatClient = chatClient;
    }

    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        AssistantMessage replySql = (AssistantMessage) state.value("reply_sql").orElse(null);
        String sql = (String) state.value("input").orElse("");
        if (replySql != null && replySql.getText() != null) {
            sql = replySql.getText();
        }
        List<TableMetadata> tableMetadataList;
        try {
            tableMetadataList = SqlParserUtil.parser(sql);
            if (CollUtil.isEmpty(tableMetadataList)) {
                tableMetadataList = chatClient.prompt()
                        .system("你是一个资深Java程序员，你擅长基于用户输入的SQL语句，解析为表结构实体类列表。")
                        .user(sql)
                        .call()
                        .entity(new ParameterizedTypeReference<List<TableMetadata>>() {
                        });
            }
        } catch (JSQLParserException e) {
            tableMetadataList = chatClient.prompt()
                    .system("你是一个资深Java程序员，你擅长基于用户输入的SQL语句，解析为表结构实体类列表。")
                    .user(sql)
                    .call()
                    .entity(new ParameterizedTypeReference<List<TableMetadata>>() {
                    });
        }
        return Map.of("llm_response", tableMetadataList);
    }
}
